#include "CHardware.HDD.h"
#include "CHardware.h"
#include <Windows.h>
#include <winioctl.h>
#include <QDebug>

#pragma warning(push)
#pragma warning(disable: 4996)

namespace NOpenHardwareMonitor {

namespace NHardware {
	
const WORD IDE_ATAPI_IDENTIFY = 0xA1;
const WORD IDE_ATA_IDENTIFY = 0xEC;


class CHDDPrivate
{
public:
	CHardware* pHardware = nullptr;


public:
	QString mIdentifier;
	QString mSerialNumber;
	QString mModelNumber;
	QString mFirmwareNumber;


public:
	BOOL Trim(char* szStr);
	char* ConvertToString(DWORD dwDiskData[256], int iFirstIndex, int iLastIndex);
	BOOL DoIdentify(HANDLE hPhysicalDriveIOCTL, PSENDCMDINPARAMS pSCIP, PSENDCMDOUTPARAMS pSCOP, BYTE btIDCmd, BYTE btDriveNum, PDWORD pdwBytesReturned);
	int GetDiskInfo(int iDriver, char* szSerialNumber, char* szModelNumber, char* szFirmwareNumber);
};

BOOL CHDDPrivate::Trim(char* szStr)
{
	int i = 0, j = 0, iFirst = -1, iLast = -1;
	int iLen = strlen(szStr);
	char szTemp[256] = { 0 };

	for (i = 0; i < iLen; i++)
	{
		if (' ' != szStr[i])
		{
			iFirst = i;
			break;
		}
	}

	for (i = (iLen - 1); 0 <= i; i--)
	{
		if (' ' != szStr[i])
		{
			iLast = i;
			break;
		}
	}

	if (-1 == iFirst || -1 == iLast)
	{
		return FALSE;
	}

	for (i = iFirst; i <= iLast; i++)
	{
		szTemp[j] = szStr[i];
		j++;
	}
	szTemp[j] = '\0';
	strcpy(szStr, szTemp);

	return TRUE;
}

char* CHDDPrivate::ConvertToString(DWORD dwDiskData[256], int iFirstIndex, int iLastIndex)
{
	static char szResBuf[256];
	int iIndex = 0;
	int iPosition = 0;

	for (iIndex = iFirstIndex; iIndex <= iLastIndex; iIndex++)
	{
		szResBuf[iPosition] = (char)(dwDiskData[iIndex] / 256);
		iPosition++;

		szResBuf[iPosition] = (char)(dwDiskData[iIndex] % 256);
		iPosition++;
	}
	szResBuf[iPosition] = '\0';

	Trim(szResBuf);
	return szResBuf;
}
BOOL CHDDPrivate::DoIdentify(HANDLE hPhysicalDriveIOCTL, PSENDCMDINPARAMS pSCIP, PSENDCMDOUTPARAMS pSCOP, BYTE btIDCmd, BYTE btDriveNum, PDWORD pdwBytesReturned)
{
	pSCIP->cBufferSize = IDENTIFY_BUFFER_SIZE;
	pSCIP->irDriveRegs.bFeaturesReg = 0;
	pSCIP->irDriveRegs.bSectorCountReg = 1;
	pSCIP->irDriveRegs.bSectorNumberReg = 1;
	pSCIP->irDriveRegs.bCylLowReg = 0;
	pSCIP->irDriveRegs.bCylHighReg = 0;
	pSCIP->irDriveRegs.bDriveHeadReg = (btDriveNum & 1) ? 0xB0 : 0xA0;
	pSCIP->irDriveRegs.bCommandReg = btIDCmd;
	pSCIP->bDriveNumber = btDriveNum;

	return DeviceIoControl(hPhysicalDriveIOCTL, SMART_RCV_DRIVE_DATA, (LPVOID)pSCIP, sizeof(SENDCMDINPARAMS) - 1,
		(LPVOID)pSCOP, sizeof(SENDCMDOUTPARAMS) + IDENTIFY_BUFFER_SIZE - 1, pdwBytesReturned, NULL);
	return FALSE;
}

int CHDDPrivate::GetDiskInfo(int iDriver, char* szSerialNumber, char* szModelNumber, char* szFirmwareNumber)
{
	char szFilePath[64] = { 0 };
	sprintf(szFilePath, "\\\\.\\PHYSICALDRIVE%d", iDriver);

	HANDLE hFile = CreateFileA(szFilePath, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
	if (INVALID_HANDLE_VALUE == hFile)
	{
		return -1;
	}

	DWORD dwBytesReturned = 0;
	GETVERSIONINPARAMS gvopVersionParam;
	DeviceIoControl(hFile, SMART_GET_VERSION, NULL, 0, &gvopVersionParam, sizeof(gvopVersionParam), &dwBytesReturned, NULL);
	if (0 >= gvopVersionParam.bIDEDeviceMap)
	{
		return -2;
	}

	unsigned int uiIDCmd = 0;
	SENDCMDINPARAMS InParams;
	unsigned int uiDrive = 0;
	uiIDCmd = (gvopVersionParam.bIDEDeviceMap >> uiDrive & 0x10) ? IDE_ATAPI_IDENTIFY : IDE_ATA_IDENTIFY;

	BYTE btOutCmd[sizeof(SENDCMDOUTPARAMS) + IDENTIFY_BUFFER_SIZE - 1];
	if (FALSE == DoIdentify(hFile, &InParams, (SENDCMDOUTPARAMS*)btOutCmd, (BYTE)uiIDCmd, (BYTE)uiDrive, &dwBytesReturned))
	{
		return -3;
	}

	CloseHandle(hFile);

	DWORD dwDiskData[256];
	USHORT* pIDSector = NULL;

	pIDSector = (USHORT*)((SENDCMDOUTPARAMS*)btOutCmd)->bBuffer;
	for (int i = 0; i < 256; i++)
	{
		dwDiskData[i] = pIDSector[i];
	}

	strcpy(szSerialNumber, ConvertToString(dwDiskData, 10, 19));
	strcpy(szModelNumber, ConvertToString(dwDiskData, 27, 46));
	strcpy(szFirmwareNumber, ConvertToString(dwDiskData, 23, 26));

	return 0;
}





CHDD::CHDD(CHardware* Hardware, const QString& HDDIdentifier)
	: QObject(Hardware)
	, d_ptr(new CHDDPrivate())
{
	Q_D(CHDD);
	d->pHardware = Hardware;
	d->mIdentifier = HDDIdentifier;

	QStringList m_IdentifierSplit = HDDIdentifier.split("/");
	if (m_IdentifierSplit.length() > 0)
	{
		qint32 m_Identifier = m_IdentifierSplit[m_IdentifierSplit.length() - 1].toInt();

		char m_SerialNumber[256];
		char m_ModelNumber[256];
		char m_FirmwareNumber[256];
		if (0 == d->GetDiskInfo(m_Identifier, m_SerialNumber, m_ModelNumber, m_FirmwareNumber))
		{
			d->mSerialNumber = m_SerialNumber;
			d->mModelNumber = m_ModelNumber;
			d->mFirmwareNumber = m_FirmwareNumber;
		}
	}
}

CHDD::~CHDD()
{  }

void CHDD::Release()
{ delete this; }

QString CHDD::GetSerialNumber() const
{ Q_D(const CHDD); return d->mSerialNumber; }

QString CHDD::GetModelNumber() const
{ Q_D(const CHDD); return d->mModelNumber; }

QString CHDD::GetFirmwareNumber() const
{ Q_D(const CHDD); return d->mFirmwareNumber; }

QString CHDD::GetIdentifier() const
{ Q_D(const CHDD); return d->mIdentifier; }


}}

#pragma warning(pop)
